from core.function import Function
import utils

class Col2im(Function):
    def __init__(self, input_shape, kernel_size, stride, pad, to_matrix):
        """
        Col2im 操作的函数类。

        Args:
            input_shape (tuple): 输入数据的形状。
            kernel_size (int or tuple): 池化窗口的大小，可以是单个整数或包含两个整数的元组。
            stride (int or tuple): 池化操作的步幅，可以是单个整数或包含两个整数的元组。
            pad (int or tuple): 池化操作的填充，可以是单个整数或包含两个整数的元组。
            to_matrix (bool): 是否将结果转换为矩阵形式。
        """
        super().__init__()
        self.input_shape = input_shape
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad
        self.to_matrix = to_matrix

    def forward(self, x):
        """
        前向传播，执行 Col2im 操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: Col2im 操作的输出变量。
        """
        y = utils.functions_conv.col2im_array(x, self.input_shape, self.kernel_size, self.stride,
                         self.pad, self.to_matrix)
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的梯度。
        """
        from utils.functions_collect import im2col
        gx = im2col(gy, self.kernel_size, self.stride, self.pad,
                    self.to_matrix)
        return gx
